# -*- coding: utf-8 -*-
import asyncio
import json
import re
from uuid import uuid4

from fastapi import FastAPI, HTTPException
from pathlib import Path

from langchain_core.documents import Document
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel
from fastapi.responses import StreamingResponse
from fastapi.concurrency import run_in_threadpool
import os
import pandas as pd
from fastapi.middleware.cors import CORSMiddleware

from utils.scheduled_tasks import scheduler
from utils.log import logger_all, logger_ass, logger_err
from utils.tools import llm, config, pre_question_search, rerank
from utils.weaviate_hybrid_search import WeaviateHybridRetriever

zs = FastAPI()
zs.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# api启动前执行
@zs.on_event("startup")
async def startup():
    global rag_chain
    scheduler.start()
    print("index、model加载中...")
    #
    # bge_rerank = FlagLLMReranker(
    #     top_n=config["llm_rerank"]["top_k"],
    #     model=config["llm_rerank"]["model_name"],
    #     device=config["llm_rerank"]["device"]
    # )

    message = """\
    ## Role :
    你是一个值得信赖的专家问答系统。
    ## Goals :
    1. 你会始终使用提供的上下文信息回答查询，而不是先前的知识。
    2. 若没有获得提供的上下文信息，你会礼貌的说明原因并表达无法回答。
    ## Constrains :
    1. 不要在你的回答中直接引用给定的上下文。
    2. 避免使用诸如'根据上下文,…'或'上下文信息…'等类似语句。
    上下文信息如下。每个上下文都以一个参考数字开头，比如[[citation:x]]，其中x是一个数字。请使用上下文，并在每句话的末尾引用上下文（如适用）。
    格式为[citation:x]。如果一个句子来自多个上下文，请列出所有适用的引文，如[citation:3][citation:5]。
    ---------------------
    {context}
    ---------------------
    回答以下查询，注意是根据上下文信息而不是先前的知识，否则你会受到惩罚。若没有获得提供的上下文信息，你会礼貌的说明原因并表达无法回答。
    查询: {query}
    答案: 
    """
    prompt = ChatPromptTemplate.from_messages([("system", "你是云瀚联盟的知识问答系统。"), ("human", message)])

    rag_chain = prompt | llm

@zs.on_event("shutdown")
async def shutdown_event():
    # 停止调度器
    scheduler.shutdown()


_more_questions_prompt = """
你是一个智能且高效的推荐问题助手。

以下是问题的上下文信息：
--------------------
{context}
--------------------

## 工作流程
1. 逐步思考并理解问题的上下文信息，根据问题的上下文信息，生成十个不同的问答对。
2. 理解用户输入的问题，若用户问题与提供的上下文信息不相关时，只输出“无推荐问题”，不生成任何推荐问题。
3. 若用户问题与提供的上下文信息相关时，逐步思考并评估步骤1中生成的问答对与用户原始问题的相关度，删除生成的十个问答对中的问题与用户原始问题<相同或高度相似>的问题，比如“上海同济大学是不是联盟成员？”和“上海同济大学是联盟成员吗？”。
4. 根据步骤3中剩余问答对中的问题与用户原始问题的相关度评估结果，从剩余问答对中选择<最不相关>的三个问答对中的问题作为推荐问题。
5. 务必<直接输出>步骤4中生成的三个推荐问题，输出之前确保没有“-”、序号等多余字符。问题之间仅使用换行符分割。

以下是用户输入的问题：
--------------------
{query}
--------------------

请使用中文回答。通过<工作流程>根据用户问题和问题的上下文信息生成三个推荐问题。输出前务必确保只把最终的<三个推荐问题>提取出来，不要生产其他内容。若没有找到提供的上下文信息只输出“无推荐问题”，不生成任何推荐问题，否则你会受到惩罚。
"""

# 推荐问题
def _create_query(query: str, context: list[Document]):
    # 拼接背景知识
    background = ""
    for item in context:
        background += f"{item.page_content} \n"

    res = llm.invoke([HumanMessage(content=_more_questions_prompt.format(context=background, query=query))])
    return [s for s in res.content.split("\n") if len(s.strip())>1]

def _generate(query: str):
    try:
        uuid = str(uuid4())
        logger_all.info({"id": uuid, "query": query})

        def stream_content():
            # 先检索预设问题
            is_pre = False
            pre_answer = pre_question_search(
                collection=config["weaviate"]["pre_query_collection"]["name"],
                query=query,
                certainty=0.95,
                limit=1
            )
            if len(pre_answer)>0:
                logger_all.info({"id": uuid, "answer": pre_answer[0].page_content})
                for answer in pre_answer[0].page_content:
                    yield json.dumps({"code": 200, "data": {"content": answer, "done": False}, "type": "summary"}) + "\n"
                yield json.dumps({"code": 200, "data": {"content": "", "done": True}, "type": "summary"}) + "\n"
                is_pre = True
            else:
                pre_answer = pre_question_search(
                    collection=config["weaviate"]["pre_query_collection"]["name"],
                    query=query,
                    certainty=0.8,
                    limit=10
                )
            #检索知识库
            retriever = WeaviateHybridRetriever(top_k=40)
            source_nodes = retriever.invoke(query)
            source_nodes.extend(pre_answer)
            # score_nodes = bge_rerank.postprocess_nodes(source_nodes, query, 0)
            score_nodes = rerank(query, source_nodes, 0)

            def format_docs(docs):
                return "\n\n".join(doc.page_content for doc in docs)
            if not is_pre:
                mid = ""
                for cr in rag_chain.stream({"context": format_docs(score_nodes), "query": query}):
                    mid += cr.content
                    yield json.dumps({"code": 200, "data": {"content": cr.content, "done": False}, "type": "summary"}) + "\n"
                yield json.dumps({"code": 200, "data": {"content": "", "done": True}, "type": "summary"}) + "\n"
                logger_all.info({"id": uuid, "answer": mid})
            source_urls = []
            source_nodes = []

            for node in score_nodes:
                if node.metadata["title"] != "预设问题":
                    source_urls.append({"url": node.metadata['url'], "title": node.metadata['title']})
                source_nodes.append({"text": node.page_content, "metadata": node.metadata})

            logger_all.info({"id": uuid, "retrievers": source_nodes})
            df = pd.DataFrame(source_urls)
            unique_df = df.drop_duplicates(subset=["url"])
            unique_urls = unique_df.to_dict("records")
            for item in unique_urls:
                item["url"] = None if item["url"] == item["title"] else item["url"]

            # 生成推荐问题
            query_list = _create_query(query, score_nodes)
            for q in query_list:
                match = re.match(r'^\d+[.)\s]*', q)
                if match:
                    q = re.sub(r'^\d+[.)\s]*', '', q)
            logger_all.info({"id": uuid, "recommend": query_list})
            if len(query_list) == 1:
                return {"end"}
            yield json.dumps({"code": 200, "data": {"source_urls": unique_urls}, "type": "source"}) + "\n"
            yield json.dumps({"code": 200, "data": {"queries": query_list}, "type": "recommend"})


        return StreamingResponse(
            stream_content(),
            headers={'Content-Type': 'application/x-ndjson'},
        )

    except Exception as e:
        logger_err.error(e)
        return {"code": 500, "error": "connect llm error"}


@zs.post("/zs/query", description="问答")
async def assistant_robot(query: str):
    return await run_in_threadpool(_generate, query)


class Query(BaseModel):
    rating: int = None
    describe: str = None
    timestamp: float = None
    query: str = None
    answer: str = None


@zs.put("/zs/feedback", description="结果反馈")
async def feedback(query: Query):
    try:
        logger_ass.info(json.dumps(query.dict()))
        return {"code": 200, "msg": "success"}
    except Exception as e:
        print(e)
        logger_err.error(e)
        return {"code": 500, "msg": "error"}


from fastapi.staticfiles import StaticFiles
from starlette.exceptions import HTTPException as StarletteHTTPException
class SPAStaticFiles(StaticFiles):
    async def get_response(self, path: str, scope):
        try:
            return await super().get_response(path, scope)
        except (HTTPException, StarletteHTTPException) as ex:
            if ex.status_code == 404:
                return await super().get_response("index.html", scope)
            else:
                raise ex


FRONTEND_BUILD_DIR = str(Path(os.getenv("FRONTEND_BUILD_DIR", "./dist")))
zs.mount(
    "/",
    SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
    name="spa-static-files",
)

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(zs, host="0.0.0.0", port=8025, loop="asyncio")
